import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from tftb.processing.cohen import WignerVilleDistribution, PageRepresentation, MargenauHillDistribution
from scipy.signal import hilbert, spectrogram, get_window
import numpy as np
from numpy import log as ln
from datetime import date, timedelta
import tqdm


def find_max_min(arr):
    if max(arr) > abs(min(arr)):
        return arr / max(arr), max(arr / max(arr))

    else:
        return arr / abs(min(arr)), abs(min(arr / abs(min(arr))))

def elipse(kx, ky, kz):
    u = np.linspace(0, 2 * np.pi, 20)
    v = np.linspace(0, np.pi, 20)

    return kx * np.outer(np.cos(u), np.sin(v)), ky * np.outer(np.sin(u), np.sin(v)), kz * np.outer(np.ones(np.size(u)), np.cos(v))

def Magnetic_3d_picture(x, y, z, begin, end, step):


    for i in range(1, step + 1):
        x1, kx = find_max_min(x[begin + (i - 1) * 720:begin + i * 720])
        y1, ky = find_max_min(y[begin + (i - 1) * 720:begin + i * 720])
        z1, kz = find_max_min(z[begin + (i - 1) * 720:begin + i * 720])

        x2, y2, z2 = elipse(kx, ky, kz)

        fig = plt.figure(figsize=(10, 10))
        ax = fig.add_subplot(projection='3d')
        ax.scatter(x1, y1, z1, c='black', marker=">")
        ax.plot_wireframe(x2, y2, z2, linewidth=0.1, color='b')
        date_begin = date(2023, 1, 1) + timedelta(days=df_B.iloc[begin + (i - 1) * 720]["Day"])
        plt.title("Day: " + str(date_begin) + " Hour: " + str(df_B.iloc[begin + (i - 1) * 720]["Hour"]) + " - " + str(
            df_B.iloc[begin + (i) * 720]["Hour"]))
        plt.xlabel('$Bx$', color='g')
        plt.ylabel('$By$', color='g')
        ax.set_zlabel('$Bz$')
        # plt.savefig(f"/home/kmg/git/TFTB_GUIproject/Figure/diagram_{i}.png", dpi=300)
        plt.show()


def combination_rotation_animation(x, y, z, begin, amoung, step):
    from matplotlib.animation import PillowWriter, FFMpegWriter

    metadata = dict(title="Movie")
    writer1 = PillowWriter(fps=15, metadata=metadata)
    writer2 = FFMpegWriter(fps=15, metadata=metadata)
    for j in writer1, writer2:
        if j == writer1:
            text = ".gif"
        else:
            text = ".mp4"

        fig = plt.figure(figsize=(12, 12), dpi=100)
        ax = fig.add_subplot(projection='3d')

        interval = 360 / amoung
        count = 0
        start = 0
        total = 360

        with j.saving(fig, "B_animate"+text, 100):
            for i in tqdm.tqdm(np.linspace(-180, 180, total)):
                ax.clear()

                count += 360 / total
                if count > interval:
                    count = 0
                    start += step

                x1, kx = find_max_min(x[start: start + step])
                y1, ky = find_max_min(y[start: start + step])
                z1, kz = find_max_min(z[start: start + step])
                x2, y2, z2 = elipse(kx, ky, kz)

                ax.plot_wireframe(x2, y2, z2, linewidth=0.1, color='b')
                ax.scatter(x1, y1, z1, c='black',
                           marker=">")
                ax.view_init(azim=i, elev=28)
                date_begin = date(2023, 1, 1) + timedelta(days=df_B.iloc[begin+start]["Day"])
                plt.title(
                    "Day: " + str(date_begin) + " Hour: " + str(df_B.iloc[begin+start]["Hour"]) + " - " + str(
                        df_B.iloc[begin + start + step]["Hour"]) + " Point: " + str(begin+start) + " - " + str(begin + start + step), size=20)
                plt.xlabel('$Bx$', size=20, color='g')
                plt.ylabel('$By$', size=20, color='g')
                ax.set_zlabel('$Bz$', size=20)
                j.grab_frame()


def Velocity_3d_picture(x, y, z, begin=0, end=0):
    step = 5760
    for i in range(begin, end, step):
        x1 = x[i:i+step]
        y1 = y[i:i+step]
        z1 = z[i:i+step]
        fig = plt.figure(figsize=(10, 10))
        ax = fig.add_subplot(projection='3d')
        ax.scatter(x1, y1, z1, c='black', marker=">")
        date_begin = date(2023, 1, 1) + timedelta(days=df_V.iloc[i]["Day"]-1)
        plt.title("Day: " + str(date_begin), size=20)
        plt.xlabel('$Vx$', color='g', size=20)
        plt.ylabel('$Vy$', color='g', size=20)
        ax.set_zlabel('$Vz$', size=20)
        plt.savefig(f"/home/kmg/git/TFTB_GUIproject/Figure/diagram_{i}.png", dpi=300)
        #plt.show()

    # animation_rotate(x1,y1,z1)

def velocity_rotation(x, y, z, begin, amoung, step):
    from matplotlib.animation import PillowWriter, FFMpegWriter

    metadata = dict(title="Movie")
    writer1 = PillowWriter(fps=15, metadata=metadata)
    writer2 = FFMpegWriter(fps=15, metadata=metadata)
    for j in writer1, writer2:
        if j == writer1:
            text = ".gif"
        else:
            text = ".mp4"

        fig = plt.figure(figsize=(12, 12), dpi=100)
        ax = fig.add_subplot(projection='3d')

        interval = 360 / amoung
        count = 0
        start = 0
        total = 360

        with j.saving(fig, "V_animate" + text, 100):
            for i in tqdm.tqdm(np.linspace(-180, 180, total)):
                ax.clear()

                count += 360 / total
                if count > interval:
                    count = 0
                    start += step
                ax.scatter(x[start: start + step], y[start: start + step], z[start: start + step], c='black',
                           marker=">")
                ax.view_init(azim=i, elev=28)
                date_begin = date(2023, 1, 1) + timedelta(days=df_V.iloc[begin + start]["Day"]-1)
                plt.title("Day: " + str(date_begin), size=20)
                plt.xlabel('$Vx$', color='g', size=20)
                plt.ylabel('$Vy$', color='g', size=20)
                ax.set_zlabel('$Vz$', size=20)
                j.grab_frame()
def animation_rotate(x, y, z):
    from matplotlib.animation import PillowWriter

    metadata = dict(title="Movie")
    writer = PillowWriter(fps=15, metadata=metadata)
    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax = fig.add_subplot(projection='3d')

    with writer.saving(fig, "V_animate.gif", 100):
        for i in np.linspace(-180, 180, 1440):
            ax.clear()
            ax.scatter(x, y, z, c='black', marker=">")
            ax.view_init(azim=i, elev=28)
            plt.xlabel('$Vx$', color='g')
            plt.ylabel('$Vy$', color='g')
            ax.set_zlabel('$Vz$')
            writer.grab_frame()


def Anime(x, y, z, len_step):
    from matplotlib.animation import PillowWriter

    metadata = dict(title="Movie")
    writer = PillowWriter(fps=15, metadata=metadata)

    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax = fig.add_subplot(projection='3d')

    def update(i):
        ax.clear()
        ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step],
                   z[i * len_step:i * len_step + len_step], s=0.5)
        plt.title(str(i * len_step))
        plt.xlabel('$Bx$', color='g')
        plt.ylabel('$By$', color='g')
        ax.set_zlabel('$Bz$')

    ani = animation.FuncAnimation(fig, update, np.arange(100), interval=1000, repeat=False)

    with writer.saving(fig, "Banimate-fix.gif", 100):
        for i in range(100):
            ax.clear()
            ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step],
                       z[i * len_step:i * len_step + len_step])
            plt.title(str(i * len_step))
            plt.xlabel('$Bx$', color='g')
            plt.ylabel('$By$', color='g')
            ax.set_zlabel('$Bz$')
            writer.grab_frame()
    plt.show()


def Video(x, y, z, len_step):
    from matplotlib.animation import FFMpegWriter

    metadata = dict(title="Movie")
    writer2 = FFMpegWriter(fps=2, metadata=metadata)

    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax = fig.add_subplot(projection='3d')

    with writer2.saving(fig, "Banimate.mp4", 100):
        for i in range(100):
            ax.clear()
            ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step],
                       z[i * len_step:i * len_step + len_step])
            plt.title(str(i * len_step))
            plt.xlabel('$Bx$', color='g')
            plt.ylabel('$By$', color='g')
            ax.set_zlabel('$Bz$')
            # ax.set_xlim(min(X), max(X))
            # ax.set_ylim(min(Y), max(Y))
            # ax.set_zlim(min(Z), max(Z))
            writer2.grab_frame()
    plt.show()

    fig = plt.figure(figsize=plt.figaspect(0.5))
    amoung = 3
    j = 1
    for i in range(24, 33):
        ax = fig.add_subplot(amoung, amoung, j, projection='3d')
        ax.scatter(x[i * len_step:i * len_step + len_step], y[i * len_step:i * len_step + len_step],
                   z[i * len_step:i * len_step + len_step], s=0.01)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        ax.set_zticklabels([])
        j += 1
    plt.show()


def B_Diagrams(df, begin, end):
    X = df["Bx_GSE"]
    Y = df["By_GSE"]
    Z = df["Bz_GSE"]

    step = (end - begin) // 720 # шаг по 3 часа для 15сек

    combination_rotation_animation(X[begin:end], Y[begin:end], Z[begin:end], begin, step, 720)
    # Magnetic_3d_picture(X, Y, Z, begin, end, step)

    # Anime(X, Y, Z, step)


def V_Diagrams(begin, end):
    X = df_V["Vx_GSE"]
    Y = df_V["Vy_GSE"]
    Z = df_V["Vz_GSE"]
    step = (end - begin) // 5760
    Velocity_3d_picture(X, Y, Z, begin, end)
    velocity_rotation(X[begin:end], Y[begin:end], Z[begin:end], begin, step, 5760)

def Read_B_Table(file):
    arr = ["Year", "Day", "Hour", "Minute", "Sec", "Bx_GSE", "By_GSE", "Bz_GSE"]
    df = pd.read_table(file, names=arr, sep="\s+")
    # for i in arr:
    #    df.loc[(df[i] > 9999)] = None
    return df


def Read_V_Table(file):
    arr = ["Year", "Day", "Hour", "Minute", "Vx_GSE", "Vy_GSE", "Vz_GSE"]
    df = pd.read_table(file, names=arr, sep="\s+")
    for i in arr:
        df.loc[(df[i] > 9999)] = None
    return df


def create_table(table):
    c = 0
    max_c = 0
    index = 0
    index2 = 0
    for i in table:
        index += 1

        if i > 999:
            c = 0
        else:
            c += 1

        if c > max_c:
            index2 = index
            max_c = c

    return index2 - max_c, index2

def simple_plot(df):
    fig = plt.figure(figsize=(12, 12), dpi=100)
    ax1 = fig.add_subplot(311)
    ax1.plot(df.iloc[:,4])
    ax2 = fig.add_subplot(312)
    ax2.plot(df.iloc[:,5])
    ax3 = fig.add_subplot(313)
    ax3.plot(df.iloc[:,6])
    plt.show()

if __name__ == "__main__":
    # df_B = Read_B_Table('sun.lst')
    # simple_plot()
    # start, end = create_table(df_B["Bx_GSE"])
    # # start, end = 309000, 314760
    # #start, end = 180000, 185760
    # print(start, end)
    # B_Diagrams(df_B, start, end)


    df_V = Read_V_Table('sun_vel.lst')
    simple_plot(df_V)
    # start, end = create_table(df_V["Vx_GSE"])
    start, end = 0, len(df_V)
    print(start, end)
    V_Diagrams(start,end)
